from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import pandas as pd

driver = webdriver.Chrome(service=Service('D:\\myproject\\feng-huang\\1.自动化\\1.selenium\\chromedriver.exe'))
driver.maximize_window()
# 设置隐式等待时间
driver.implicitly_wait(20)

def extract_data(name, dict):
    try:
        title = name.find_element(By.XPATH, './a/div[2]/h2')
        dict['简称'] = title.text
    except AttributeError:
        dict['简称'] = ''

    try:
        dict['产品类别'] = name.find_element(By.XPATH, './a/div[2]/div[1]').text
    except:
        dict['产品类别'] = ''

    try:
        dict['排名'] = name.find_element(By.CLASS_NAME, 'rankings').text
    except:
        dict['排名'] = ''

    try:
        dict['位置'] = name.find_element(By.CLASS_NAME, 'location').text
    except:
        dict['位置'] = ''

    try:
        dict['产品'] = name.find_element(By.CLASS_NAME, 'products').text
    except:
        dict['产品'] = ''

    try:
        url = name.find_element(By.XPATH, './a')
        dict['url'] = url.get_attribute('href')
    except:
        dict['url'] = ''

def extract_page_data(page_num):
    try:
        data = []
        driver.get(f'https://www.foodtalks.cn/company/select?page={page_num}')
        names = driver.find_elements(By.XPATH, '/html/body/div[1]/div[2]/section/main/section/main/div[1]/div[2]/div[2]/div')
        for name in names:
            dict = {'简称': '', '产品类别': '', '排名': '', '位置': '', '产品': '', 'url': '',
                    '公司名称': '', '公司简称': '', '英文名称': '', '英文简称': '',
                    '公司类型': '', '成立时间': '', '股票代码': '', '公司总机': '',
                    '公司性质': '', '行业分类': '', '主要品牌': '', '主要产品': '',
                    '公司官网': '', '公司总部': '', '总部地址': '', '公司介绍': '',
                    '品牌': '', '客户': '', '供应商': '', '联系方式': ''
                    }
            extract_data(name, dict)
            title = name.find_element(By.XPATH, './a/div[2]/h2')
            title.click()
            WebDriverWait(driver, 10).until(EC.number_of_windows_to_be(2))
            driver.switch_to.window(driver.window_handles[1])
            try:
                driver.find_element(By.XPATH, '//span[text()="首页"]').click()
            except:
                pass

            fields = {
                '公司名称': '//div[text()="中文全称："]',
                '公司简称': '//div[text()="中文简称："]',
                '英文名称': '//div[text()="英文全称："]',
                '英文简称': '//div[text()="英文简称："]',
                '公司类型': '//div[text()="公司类型："]',
                '成立时间': '//div[text()="成立时间："]',
                '股票代码': '//div[text()="股票代码："]',
                '公司总机': '//div[text()="公司总机："]',
                '公司性质': '//div[text()="公司性质："]',
                '行业分类': '//div[text()="行业分类："]',
                '主要品牌': '//div[text()="主要品牌："]',
                '主要产品': '//div[text()="主要产品："]',
                '公司官网': '//div[text()="公司官网："]',
                '公司总部': '//div[text()="公司总部："]',
                '总部地址': '//div[text()="总部地址："]'
            }

            for key, field in fields.items():
                try:
                    element = driver.find_element(By.XPATH, field)
                    if element.text == f'{key[:-1]}：':
                        try:
                            dict[key] = element.find_element(By.XPATH, "./following-sibling::*[1]").text
                        except:
                            dict[key] = ''
                except:
                    pass

            # 简介
            try:
                unfold_more = driver.find_element(By.XPATH, '//span[text()="展开"]')
                if unfold_more:
                    unfold_more.click()
            except:
                pass

            try:
                introduction = driver.find_element(By.XPATH, '//span[text()="公司介绍"]')
                if introduction.text == '公司介绍':
                    dict['公司介绍'] = driver.find_element(By.XPATH, '//*[@id="app"]/div[2]/section/main/section/main/div[1]/div[1]/div[4]/div/div[1]/div[1]/div[2]/div[1]').text
            except:
                pass

            # 品牌、产品、客户、供应商、联系方式
            for btn_text in ['品牌', '产品', '客户', '供应商', '联系方式']:
                try:
                    btn = driver.find_element(By.XPATH, f'//span[text()="{btn_text}"]')
                    if btn.text == btn_text:
                        btn.click()
                        WebDriverWait(driver, 10).until(EC.presence_of_element_located((By.CLASS_NAME, 'el-row')))
                        elements = driver.find_elements(By.CLASS_NAME, 'el-row' if btn_text in ['品牌', '产品'] else 'customer-list.not-login')
                        dict[btn_text] = [ele.text for ele in elements]
                except:
                    dict[btn_text] = ''

            data.append(dict)
            driver.close()
            driver.switch_to.window(driver.window_handles[0])

        df = pd.DataFrame(data)
        df.to_excel(f'./data/company_data_page_{page_num}.xlsx', index=False, engine='openpyxl')
        print(f'第{page_num}页数据获取并保存完成')

    except Exception as e:
        print(f"在第{page_num}页发生错误: {e}")
        driver.quit()

for i in range(1, 485):
    extract_page_data(i)

driver.quit()
